#!/usr/bin/env python3

#==============================================================================
# Functionality
#==============================================================================
import pdb
import sys
import os
import re
import struct

# utility funcs, classes, etc go here.

def asserting(cond):
    if not cond:
        pdb.set_trace()
    assert(cond)

def has_stdin():
    return not sys.stdin.isatty()

def reg(pat, flags=0):
    return re.compile(pat, re.VERBOSE | flags)

#==============================================================================
# Cmdline
#==============================================================================
import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, 
    description="""
TODO
""")
     
parser.add_argument('-v', '--verbose',
    action="store_true",
    help="verbose output" )
     
parser.add_argument('-M', '--max-size',
    type=int,
    default=None,
    help="Skip any events whose serialized string size (before base64) is greater than this value." )

args = None

#==============================================================================
# Main
#==============================================================================

from base64 import b64encode

def parse_event(f, offset):
  header_str = f.read(8)
  if not header_str:
    return offset, None
  n = struct.unpack("<Q", header_str)[0]
  crc_header_str = f.read(4)
  event_str = f.read(n)
  crc_event_str = f.read(4)
  return offset + len(event_str) + 16, event_str

def parse_tfevents(f):
  offset = 0
  while True:
    offset, event_str = parse_event(f, offset)
    if event_str is None:
      break
    if args.max_size is None or len(event_str) <= args.max_size:
      yield offset, event_str

def parse_tfevents_stream(f):
  for offset, event_str in parse_tfevents(f):
    sys.stdout.write(str(offset) + ',' + b64encode(event_str).decode('utf8') + '\n')
    sys.stdout.flush()

def run():
    if args.verbose:
        print(args)
    if len(args.args) <= 0 and not has_stdin():
        # if there were no args and there was no input, prompt user.
        print('Enter input (press Ctrl-D when done):')
    if len(args.args) <= 0 or has_stdin():
        #indata = sys.stdin.buffer.read()
        parse_tfevents_stream(sys.stdin.buffer)
    # for each arg on cmdline...
    import tensorflow as tf
    for arg in args.args:
        with tf.io.gfile.GFile(arg, 'rb') as f:
          parse_tfevents_stream(f)

def main():
    try:
        global args
        if not args:
            args, leftovers = parser.parse_known_args()
            args.args = leftovers
        return run()
    except IOError:
        # http://stackoverflow.com/questions/15793886/how-to-avoid-a-broken-pipe-error-when-printing-a-large-amount-of-formatted-data
        try:
            sys.stdout.close()
        except IOError:
            pass
        try:
            sys.stderr.close()
        except IOError:
            pass

if __name__ == "__main__":
    main()

